# encoding=utf-8

""" 原始函数
def visual_wrapper(cache_path: str, image_path: str, pred_arr, result_name: str):
    obj_lst = ["background", "hemorrhages", "hard_exudates", "microaneurysms", "disc", "soft_exudates"]
    img_arr = show_seg(image_path, pred_arr)
    file_name = os.path.basename(image_path)

    # visualization
    visualization = os.path.join(cache_path, "{}/visualization".format(result_name))
    os.makedirs(visualization, exist_ok=True)

    # draw label
    img_arr = draw_text(img_arr)
    cv2.imwrite(os.path.join(visualization, file_name), img_arr)

    predict_array = np.zeros_like(pred_arr).astype(np.uint8)
    pred_flag_arr = np.argmax(pred_arr, axis=0)

    for j in range(predict_array.shape[0]):
        if not os.path.exists(os.path.join(cache_path, result_name, obj_lst[j])):
            os.makedirs(os.path.join(cache_path, result_name, obj_lst[j]), exist_ok=True)
        predict_array[j, :, :][pred_flag_arr == j] = 255
        cv2.imwrite(os.path.join(cache_path, result_name, obj_lst[j], file_name), predict_array[j, :, :])

    pred_color_arr = np.zeros((pred_flag_arr.shape[0], pred_flag_arr.shape[1], 3)).astype(np.uint8)
    pred_color_arr[:, :, 0][pred_flag_arr // 4 > 0] = 255
    pred_color_arr[:, :, 1][pred_flag_arr % 4 > 1] = 255
    pred_color_arr[:, :, -1][pred_flag_arr % 2 > 0] = 255

    if not os.path.exists(os.path.join(cache_path, "{}/all").format(result_name)):
        os.makedirs(os.path.join(cache_path, "{}/all".format(result_name)), exist_ok=True)
    cv2.imwrite(os.path.join(cache_path, "{}/all".format(result_name), file_name), pred_color_arr)
"""

import time

import cv2
import cv2 as cv
import numpy as np


def time_wrapper(f, *args, **kwargs):
    t1 = time.perf_counter()
    res = f(*args, **kwargs)
    t2 = time.perf_counter()
    print("{}, time: {}s".format(f, t2 - t1))
    return res


def render_result(ori_array: np.ndarray, pred_array: np.ndarray, color_list: list):
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    for i in range(1, pred_array.shape[0]):
        c, h = cv.findContours(cv2.dilate(pred_array[i], kernel, iterations=1), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        f_c = [_c for _c in c if cv.contourArea(_c) > 300]  # TODO 过滤数值按照病灶大小设定
        cv.drawContours(ori_array, f_c, -1, color_list[i], 2)
    return ori_array


def render_progress(ori_image, npy_arr):
    ori_array = cv.imread(ori_image)
    pred_array = np.load(npy_arr)
    # "background", "hemorrhages", "hard_exudates", "microaneurysms", "disc", "soft_exudates"
    color_list = [
        (0, 0, 0), (0, 0, 255), (0, 255, 0), (0, 255, 255), (255, 0, 0), (255, 0, 255)
    ]

    render = time_wrapper(render_result, ori_array, pred_array, color_list)
    cv.imwrite("result.png", render)


if __name__ == '__main__':
    render_progress("test_image.jpeg", "pred.npy")
